# This code is part of LINKEQ.
#
# (C) Copyright LINKE 2023.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
#
# -*- coding: utf-8 -*-
# @Time     : 2023/6/26 16:04
# @Author   : deviludier @ LINKE
# @File     : VQASoptimizer.py
# @Software : PyCharm

import  numpy as np
import time
import sys

class VQA_Scheduler:
    def __init__(self, learning_rate=0.1, y_tol=1e-10, maxit=1024):
        self.learning_rate = learning_rate
        self.y_tol = y_tol
        self.maxit = maxit

    def excute(self, loss, var_list):
        loss_list = []
        itr = 0
        while (1):
            if itr == 0:
                grads = self.gradients(loss, var_list)
                test_var_list = var_list.copy()
            else:
                grads = self.gradients(loss, updates)
                test_var_list = updates.copy()
            updates = []
            for var, grad in zip(test_var_list, grads):
                new_var = var - self.learning_rate * grad
                updates.append(new_var)
            y = loss(updates)
            loss_list.append(y)
            time.sleep(0.1)  # 模拟耗时操作
            show_progress_bar(itr + 1, self.maxit+1, prefix='Progress:', suffix='Complete', length=50)
            if (itr == self.maxit):
                break
            elif (itr > 1 and abs(loss_list[-2] - y) < self.y_tol):
                break
            # elif (itr > 1 and abs(loss_list[-1] - y) < self.y_tol):
            #     break
            else:
                itr += 1
        if itr < self.maxit:
            print('\nConvergence! Training completed! 🎉')
        else:
            print('\nSatisfying the max iteration! Training completed! 🎉')
        return loss_list, updates

    def gradients(self, loss, var_list):
        grads = []
        step = np.pi / 4
        for i in range(len(var_list)):
            tmp_var_list_p = var_list.copy()
            tmp_var_list_m = var_list.copy()
            tmp_var_list_p[i] = tmp_var_list_p[i] + step
            tmp_var_list_m[i] = tmp_var_list_m[i] - step
            grads.append((loss(tmp_var_list_p) - loss(tmp_var_list_m)) / 2 * step)
        return grads

def show_progress_bar(iteration, total, prefix='', suffix='', decimals=1, length=50, fill='█'):
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    sys.stdout.write('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix))
    sys.stdout.flush()

import numpy as np
import networkx as nx
from qiskit import QuantumCircuit, Aer, execute
from qiskit.quantum_info import Statevector

def exact_max_cut(G):
    nodes = list(G.nodes())
    n = len(nodes)
    max_cut_value = 0
    # 遍历所有可能的分割（跳过空集或全集）
    for i in range(1, 2**(n-1)):  # 利用对称性减少计算
        partition = []
        for j in range(n):
            if (i >> j) & 1:
                partition.append(nodes[j])
        cut_value = nx.cut_size(G, partition)
        if cut_value > max_cut_value:
            max_cut_value = cut_value
    return max_cut_value

if __name__ == "__main__":
    # 使用基础Qiskit实现的QAOA MaxCut类
    class QAOA_MaxCut_Qiskit:
        def __init__(self, graph, p=1):
            self.graph = graph
            self.n_qubits = len(graph.nodes())
            self.p = p  # QAOA层数
            self.backend = Aer.get_backend('statevector_simulator')
            
            # 手动构建MaxCut哈密顿量
            self.cost_hamiltonian = self._build_cost_hamiltonian()
            
        def _build_cost_hamiltonian(self):
            # 构建MaxCut哈密顿量: H_C = 1/2 * Σ_{(i,j)∈E} (I - Z_i Z_j)
            # 使用矩阵形式构建
            H_size = 2**self.n_qubits
            H_matrix = np.zeros((H_size, H_size))
            
            # 创建Pauli Z矩阵
            Z_mat = np.array([[1, 0], [0, -1]])
            I_mat = np.eye(2)
            
            for i, j in self.graph.edges():
                # 构建Z_i ⊗ Z_j项
                term = 1
                for k in range(self.n_qubits):
                    if k == i or k == j:
                        term = np.kron(term, Z_mat)
                    else:
                        term = np.kron(term, I_mat)
                
                # 添加项: 1/2 * (I - Z_i Z_j)
                H_matrix += 0.5 * (np.eye(H_size) - term)
            
            return H_matrix
        
        def create_qaoa_circuit(self, params):
            # 创建QAOA电路
            qc = QuantumCircuit(self.n_qubits)
            
            # 初始态 |+⟩^n
            for i in range(self.n_qubits):
                qc.h(i)
            
            # 添加p层QAOA
            for layer in range(self.p):
                gamma = params[2*layer] if len(params) > 2*layer else params[0]
                beta = params[2*layer + 1] if len(params) > 2*layer + 1 else params[1]
                
                # 成本哈密顿量层 - 对每条边应用RZZ门
                for edge in self.graph.edges():
                    i, j = edge
                    qc.rzz(2 * gamma, i, j)
                
                # 混合哈密顿量层 - 对每个量子比特应用RX门
                for i in range(self.n_qubits):
                    qc.rx(2 * beta, i)
            
            return qc
        
        def energy(self, params):
            # 使用Qiskit模拟QAOA电路并计算期望值
            qc = self.create_qaoa_circuit(params)
            
            # 模拟电路
            result = execute(qc, self.backend).result()
            statevector = result.get_statevector()
            
            # 将Statevector转换为numpy数组
            psi = statevector.data
            
            # 计算期望值 <ψ|H|ψ> = ψ† H ψ
            # 使用矩阵乘法计算
            H_psi = np.dot(self.cost_hamiltonian, psi)
            energy = np.real(np.dot(psi.conj(), H_psi))
            
            return -energy
    
    # 创建4个顶点、度为2的正则图（环图）
    G = nx.random_graphs.random_regular_graph(3,10)
    print("图结构:", list(G.edges()))
    
    # 创建QAOA MaxCut实例（使用Qiskit）
    qaoa = QAOA_MaxCut_Qiskit(G, p=10)  # 使用1层QAOA
    
    # 计算理论最小值（通过精确对角化）
    eigenvalues, _ = np.linalg.eig(qaoa.cost_hamiltonian)
    min_energy = float(-exact_max_cut(G))
    theoretical_min = min_energy
    
    # 最大切割值
    theoretical_max_cut = -min_energy
    
    print(f"MaxCut问题的理论最小能量: {theoretical_min:.6f}")
    print(f"对应的最大切割值: {theoretical_max_cut:.6f}")
    
    # 测试不同的最大迭代次数
    max_iterations = [10, 50, 100, 200, 500, 1000]
    
    # 记录结果
    results = []
    
    print("开始测试QAOA MaxCut优化性能（使用Qiskit）...")
    
    for max_iter in max_iterations:
        print(f"\n测试最大迭代次数: {max_iter}")
        
        # 初始化优化器
        opt = VQA_Scheduler(learning_rate=0.2, maxit=max_iter)
        
        # 初始化参数（对于p=1的QAOA，需要2个参数：beta和gamma）
        x = np.random.random(10*10) * 2 * np.pi
        
        # 定义损失函数
        def loss_function(params):
            return qaoa.energy(params)
        
        # 执行优化
        start_time = time.time()
        loss_history, final_params = opt.excute(loss=loss_function, var_list=x)
        end_time = time.time()
        
        # 计算最终能量
        final_energy = loss_function(final_params)
        
        # 计算与理论值的比例
        ratio = final_energy / theoretical_min
        
        # 计算误差（与理论最小值的绝对差）
        error = abs(final_energy - theoretical_min)
        
        # 计算近似的切割值（取负号，因为最小能量对应最大切割）
        approx_cut = -final_energy
        
        # 记录结果
        results.append({
            'max_iterations': max_iter,
            'initial_params': x,
            'final_params': final_params,
            'final_energy': final_energy,
            'approx_cut': approx_cut,
            'ratio': ratio,
            'error': error,
            'convergence_iterations': len(loss_history),
            'computation_time': end_time - start_time
        })
        
        print(f"初始参数: {x}")
        print(f"最终参数: {final_params}")
        print(f"最终能量: {final_energy:.6f}")
        print(f"近似切割值: {approx_cut:.6f}")
        print(f"理论最小能量: {theoretical_min:.6f}")
        print(f"理论最大切割值: {theoretical_max_cut:.6f}")
        print(f"与理论值比例: {ratio:.6f}")
        print(f"绝对误差: {error:.6f}")
        print(f"实际迭代次数: {len(loss_history)}")
        print(f"计算时间: {end_time - start_time:.2f}秒")
    
    # 将结果写入文件
    with open('opttest.txt', 'w') as f:
        f.write("QAOA MaxCut优化性能测试结果（使用Qiskit）\n")
        f.write("=" * 80 + "\n")
        f.write(f"图结构: {list(G.edges())}\n")
        f.write(f"理论最小能量: {theoretical_min:.6f}\n")
        f.write(f"理论最大切割值: {theoretical_max_cut:.6f}\n")
        f.write("=" * 80 + "\n")
        f.write("最大迭代次数\t最终能量\t近似切割值\t与理论值比例\t绝对误差\t实际迭代次数\t计算时间(秒)\n")
        f.write("-" * 120 + "\n")
        
        for result in results:
            f.write(f"{result['max_iterations']:>12}\t{result['final_energy']:>10.6f}\t{result['approx_cut']:>12.6f}\t"
                   f"{result['ratio']:>12.6f}\t{result['error']:>10.6f}\t{result['convergence_iterations']:>12}\t{result['computation_time']:>10.2f}\n")
    
    print(f"\n测试完成! 结果已保存到opttest.txt")
    
    # 打印简要总结
    print("\nQAOA MaxCut优化性能测试总结（使用Qiskit）:")
    print("最大迭代次数 -> 最终能量 (近似切割值)")
    for result in results:
        print(f"{result['max_iterations']:>12} -> {result['final_energy']:.6f} ({result['approx_cut']:.6f})")
    
    # 找到最佳结果（能量最低）
    best_result = min(results, key=lambda x: x['final_energy'])
    print(f"\n最佳结果: 最大迭代次数={best_result['max_iterations']}, "
          f"最终能量={best_result['final_energy']:.6f}, "
          f"近似切割值={best_result['approx_cut']:.6f}")
    
    # 检查是否找到了接近理论最小值的解
    if best_result['error'] < 0.1:  # 误差小于0.1
        print("✅ QAOA优化成功找到了接近理论最小值的解!")
        print(f"找到的切割值: {best_result['approx_cut']:.6f}, 理论最大切割值: {theoretical_max_cut:.6f}")
    else:
        print("⚠️ QAOA优化未能找到足够接近理论最小值的解，可能需要增加QAOA层数或调整优化器参数。")
    
    # 可视化最佳结果的QAOA电路
    if len(results) > 0:
        best_params = best_result['final_params']
        best_circuit = qaoa.create_qaoa_circuit(best_params)
        print(f"\n最佳QAOA电路（{qaoa.p}层）:")
        print(best_circuit.draw())